#thanks to https://www.joshuagubler.com/r-code for this code 
#TwowayME.f function slightly modified 

###############################################
###############################################
robustse.f <- function(model, cluster, df_correction) {

  require(sandwich)
  require(lmtest)
  require(multiwayvcov)
  if(missing(cluster)) {
    name <- deparse(substitute(model))
    modelname <- paste(name,"rob",sep=".")
    model$se <- coeftest(model, vcov=vcovHC(model,"HC1"))[,2]
    model$vcovHC <- vcovHC(model,"HC1")
    assign(modelname,model,envir = .GlobalEnv)
    coeftest(model, vcov=vcovHC(model,"HC1"))
  } else {
    name <- deparse(substitute(model))
    modelname <- paste(name,"rob",sep=".")
    vcovCL <- cluster.vcov(model, cluster, df_correction = df_correction)
    model$vcovCL <- vcovCL
    modelname <- paste(name,"clustrob",sep=".")
    model$se <- coeftest(model, vcovCL)[,2]
    assign(modelname,model,envir = .GlobalEnv)
    coeftest(model, vcovCL)
  }
}



###############################################
###############################################
TwowayME.f <- function(M,X,Z,xlab,ylab,level,rob,cluster,df_correction,hist,Low,High)
{
   require(multiwayvcov)
  
  if(length(na.omit(unique(Z)))>2){
    S <- summary(M)
    N <- c(1:20)
    
    ## 20 equally-spaced values for the moderating variable
    zmin <- rep(min(Z, na.rm = TRUE), 20)
    zmax <- rep(max(Z, na.rm = TRUE), 20)
    Znew <- (((N - 1) / (20 - 1)) * (zmax - zmin)) + zmin
    
    ## Grab elements of coefficient and vcov matrix
    H <- head(S$coefficients, 3)
    T <- tail(S$coefficients, 1)
    b <- rbind(H, T)
    if(missing(cluster)){
      if(missing(rob)){Vcov <- vcov(M)}
      else{
        Vcov <- vcovHC(M,"HC1")
      }
    } else{
      Vcov <- cluster.vcov(M, cluster, df_correction = df_correction)
    }
    Vcov <- as.data.frame(Vcov)
    Vcov1 <- Vcov[, c(1:3)]
    Vcov2 <- Vcov[, -c(0:0 - length(Vcov))]
    Vcov <- cbind(Vcov1, Vcov2)
    Vh <- head(Vcov, 3)
    Vt <- tail(Vcov, 1)
    V <- rbind(Vh, Vt)
    b1 <- b[2, 1]
    b3 <- b[4, 1]
    varb1 <- V[2, 2]
    varb3 <- V[4, 4]
    covb1b3 <- V[4, 2]
    
    ## Calculate ME values
    conb <- b1 + b3 * Znew
    
    ## Calculate standard errors
    conse <- sqrt(varb1 + varb3 * (Znew^2) + 2 * covb1b3 * Znew)
    
    ## Upper and lower CIs
    ci <- NA
    ci[level == 95] <- qt(.975,M$df.residual)
    ci[level == 90] <- qt(.95,M$df.residual)
    ci[level == 99] <- qt(.99,M$df.residual)

    a = ci * conse
    upper = conb + a
    lower = conb - a
    
    ##Graph the results
    if(missing(hist)){
      plot(c(Znew,Znew,Znew), c(a,upper,lower), type="n",xlab=xlab,ylab=ylab)+rug(jitter(Z))
      lines(Znew,conb,col="black")
      lines(Znew,upper,col="black",lty=2)
      lines(Znew,lower,col="black",lty=2)
      abline(h=0)
    }else{
      hist(Z, axes=F, xlab="", ylab="",main="", col="light gray")
      par(new=TRUE)
      plot(c(Znew,Znew,Znew), c(a,upper,lower), type="n",xlab=xlab,ylab=ylab)
      lines(Znew,conb,col="black")
      lines(Znew,upper,col="black",lty=2)
      lines(Znew,lower,col="black",lty=2)
      abline(h=0)  
    }   
  }else{
    S <- summary(M)
    
    ##To create the high/low dichotomous variable
    Z0 <- quantile(Z,   0, na.rm=TRUE) 
    Z1 <- quantile(Z,   1, na.rm=TRUE)
    
    ## Grab elements of coefficient and vcov matrix
    require(sandwich)
    H <- head(S$coefficients,3)
    T <- tail(S$coefficients,1)
    b <- rbind(H,T)
    if(missing(cluster)){
      if(missing(rob)){Vcov <- vcov(M)}
      else{
        Vcov <- vcovHC(M,"HC1")
      }
    } else{
      Vcov <- cluster.vcov(M, cluster, df_correction = df_correction)
    }
    Vcov <- as.data.frame(Vcov)
    Vcov1 <- Vcov[,c(1:3)]
    Vcov2 <- Vcov[,-c(0:0-length(Vcov))]
    Vcov <- cbind(Vcov1,Vcov2)
    Vh <- head(Vcov,3)
    Vt <- tail(Vcov,1)
    V <- rbind(Vh,Vt)
    b1 <- b[2,1]
    b3 <- b[4,1]
    varb1 <- V[2,2]
    varb3 <- V[4,4]
    covb1b3 <- V[4,2]
    
    ## Calculate ME values
    conb0 <- b1+b3*Z0
    conb1 <- b1+b3*Z1
    
    ## Calculate standard errors    
    conse0 <- sqrt(varb1 + varb3*(Z0^2) + 2*covb1b3*Z0)
    conse1 <- sqrt(varb1 + varb3*(Z1^2) + 2*covb1b3*Z1)
    
    ## Upper and lower CIs
    ci <- NA
    ci[level == 95] <- qt(.975,M$df.residual)
    ci[level == 90] <- qt(.95,M$df.residual)
    ci[level == 99] <- qt(.99,M$df.residual)

    a0 = ci*conse0
    a1 = ci*conse1

    upper0 = conb0 + a0
    lower0 = conb0 - a0
    upper1 = conb1 + a1
    lower1 = conb1 - a1
  

    ##Graph the results
    require(ggplot2)
    margplot <- qplot(
      c(Low,High),
      c(conb0,conb1), 
      xlab=xlab,
      ylab=ylab
    ) + 
      geom_point() +
      theme_bw()
    
    margplot + geom_linerange(aes(ymin=c(lower0,lower1),ymax=c(upper0,upper1))) + geom_abline(intercept=0,slope=0)
  }
}

